import torch
import torch.nn as nn

from spdnet.spd import SPDTransform, SPDRectified, SPDTangentSpace, SPDPower


class SPDRes2Block(nn.Module):
    """ SPD流形上的Res2Net模块 """

    def __init__(self, input_size, output_size, scal):
        super(SPDRes2Block, self).__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.scal = scal

        self.pre_transs = nn.ModuleList([
            SPDTransform(input_size, output_size) for _ in range(scal)
        ])

        self.transs = nn.ModuleList([
            SPDTransform(input_size, output_size) for _ in range(scal-1)
        ])
        # self.rects = nn.ModuleList([
        #     SPDRectified() for _ in range(scal-1)
        # ])
        # self.tangents = nn.ModuleList([
        #     SPDTangentSpace(output_size) for _ in range(scal-1)
        # ])
        self.rect = SPDRectified()
        self.tangent = SPDTangentSpace(output_size, vectorize=False)
        self.powers = nn.ModuleList([
            SPDPower(output_size) for _ in range(scal)
        ])

        self.alpha = nn.Parameter(torch.ones(scal))

    def forward(self, x):
        xs = [pre_trans(x) for pre_trans in self.pre_transs]
        # xs = [x for _ in range(self.scal)]
        ys = []
        y_prev = None
        for i, (x_i, trans, power) in enumerate(zip(xs, [None]+list(self.transs), [None]+list(self.powers[:self.scal-1]))):
            if i == 0:  # 第一个组直接传递
                y_i = x_i
            else:
                log_term = self.tangent(self.rect(trans(x_i)))
                if y_prev is not None:
                    y_i = power(self.tangent(y_prev) + log_term)
                else:
                    y_i = power(log_term)
                y_prev = log_term
            ys.append(y_i)

        log_terms = [self.tangent(y) * a for y, a in zip(ys, self.alpha)]
        z = sum(log_terms)
        out = self.powers[self.scal-1](z)

        return out



# 示例用法
if __name__ == "__main__":
    input_size = 8  # 输入SPD矩阵维度 64x64
    batch_size = 1
    X = torch.randn(batch_size, input_size, input_size)
    X = X @ X.transpose(-1, -2)  # 构造SPD矩阵

    model = SPDRes2Block(input_size, output_size=input_size, scal=4)
    Y = model(X)
    print("Input shape:", X.shape)  # [32, 64, 64]
    print("Output shape:", Y.shape)  # [32, 64, 64]
    # print(X)
    # print(Y)
    print(model)
    for name, parameter in model.named_parameters():
        if name.startswith('pre'):
            print(name)
        # print(name)
        # print(type(name))